#!/usr/bin/env python3
"""
Population ADAM PBPK Modeling Script

Simulates PBPK modeling for a population of 100 individuals using the ADAM model
(Advanced Dissolution, Absorption, and Metabolism) with population variability.

Author: PBPK Modeling System
Version: 1.0.0
Date: 2025-07-24
"""

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import json
import os
import sys
from typing import Dict, Any, List
import warnings
warnings.filterwarnings('ignore')

# Handle numpy trapezoid function compatibility
try:
    from numpy import trapezoid
except ImportError:
    # Fallback for older numpy versions
    from numpy import trapz as trapezoid

# Add path to access modules
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))

from src.absorption.absorption_module import AbsorptionModel
from src.absorption.adam_model import ADAMModel
from src.population.population_data import PopulationData
from src.visualization.plot_generator import PlotGenerator

class PopulationADAMPBPK:
    """Population-based ADAM PBPK model with comprehensive drug modeling"""
    
    def __init__(self):
        """Initialize population ADAM model"""
        self.population_size = 100
        self.base_drug_params = {
            'drug_name': 'Test_Drug',
            'dose': 100,  # mg
            'solubility': 0.8,  # mg/mL
            'permeability': 2e-4,  # cm/s
            'molecular_weight': 325.4,  # g/mol
            'pKa': 7.8,
            'logP': 2.5,
            'particle_size': 50,  # μm
            'density': 1.2,  # g/cm³
            'use_adam': True
        }
        self.population_data = PopulationData()
        self.simulation_results = {}
        
    def generate_population_parameters(self, population_type: str = "adult") -> List[Dict]:
        """Generate population with ADAM-specific parameters"""
        np.random.seed(42)  # For reproducibility
        
        # Get population parameters
        population_individuals = self.population_data.get_population_parameters(
            population_type, self.population_size
        )
        
        # Add drug-specific variability
        enhanced_population = []
        for i, individual in enumerate(population_individuals):
            # Base drug parameters with individual variability
            drug_params = self.base_drug_params.copy()
            
            # Adjust for body weight
            weight_factor = individual.get('weight', 70) / 70
            
            # Add pharmacokinetic variability
            drug_params.update({
                'id': f'individual_{i+1:03d}',
                'weight': individual.get('weight', 70),
                'age': individual.get('age', 35),
                'height': individual.get('height', 170),
                'bmi': individual.get('bmi', 24),
                'gender': 'male' if np.random.random() < 0.5 else 'female',
                'dose': drug_params['dose'] * weight_factor,
                'solubility': np.random.normal(drug_params['solubility'], 0.1),
                'permeability': np.random.lognormal(np.log(drug_params['permeability']), 0.3),
                'molecular_weight': drug_params['molecular_weight'],
                'pKa': np.random.normal(drug_params['pKa'], 0.2),
                'logP': np.random.normal(drug_params['logP'], 0.3),
                'particle_size': np.random.normal(drug_params['particle_size'], 10),
                'use_adam': True
            })
            
            # Add physiological parameters
            drug_params.update({
                'blood_volume': individual.get('blood_volume', 5.0),
                'plasma_volume': individual.get('plasma_volume', 3.0),
                'gfr': individual.get('gfr', 120),
                'cardiac_output': individual.get('cardiac_output', 6.5),
                'liver_blood_flow': individual.get('liver_blood_flow', 1.35),
                'kidney_blood_flow': individual.get('kidney_blood_flow', 1.24),
                'albumin': individual.get('albumin', 45),
                'alpha1_acid_glycoprotein': individual.get('alpha1_acid_glycoprotein', 0.8)
            })
            
            # Ensure parameter bounds
            drug_params['solubility'] = max(0.1, min(2.0, drug_params['solubility']))
            drug_params['permeability'] = max(1e-5, min(1e-3, drug_params['permeability']))
            drug_params['pKa'] = max(5.0, min(10.0, drug_params['pKa']))
            drug_params['logP'] = max(-2.0, min(5.0, drug_params['logP']))
            drug_params['particle_size'] = max(1, min(200, drug_params['particle_size']))
            
            enhanced_population.append(drug_params)
        
        return enhanced_population
    
    def run_individual_adam_simulation(self, individual_params: Dict, time_points: np.ndarray) -> Dict[str, Any]:
        """Run ADAM simulation for a single individual"""
        try:
            # Create ADAM model
            adam_model = ADAMModel(individual_params)
            
            # Create Absorption model for BCS classification
            absorption_model = AbsorptionModel(individual_params)
            
            # Calculate BCS classification
            bcs_result = absorption_model.calculate_bcs_classification(
                individual_params['solubility'],
                individual_params['permeability'],
                individual_params['dose']
            )
            
            # Calculate ADAM absorption
            adam_result = adam_model.calculate_adam_absorption(time_points)
            
            # Calculate integrated absorption
            absorption_result = absorption_model.calculate_absorption({}, time_points)
            
            # Calculate plasma concentration using proper PK modeling
            dose = individual_params['dose']
            bioavailability = adam_result['bioavailability']
            
            # Weight-adjusted PK parameters
            ke = 0.15 * (70 / individual_params.get('weight', 70)) ** 0.25  # Allometric scaling
            Vd = 50 * (individual_params.get('weight', 70) / 70)  # Weight-adjusted volume
            ka = 1.0  # Absorption rate constant (1/h)
            
            # First-order absorption model with ADAM bioavailability
            F_dose = dose * bioavailability
            
            # Calculate concentration using Bateman equation
            if abs(ka - ke) > 1e-6:
                concentration = (F_dose * ka) / (Vd * (ka - ke)) * (
                    np.exp(-ke * time_points) - np.exp(-ka * time_points)
                )
            else:
                concentration = (F_dose * ka * time_points * np.exp(-ke * time_points)) / Vd
            
            # Ensure no negative concentrations
            concentration = np.maximum(concentration, 0)
            
            # Calculate PK metrics
            cmax = np.max(concentration)
            tmax_idx = np.argmax(concentration)
            tmax = time_points[tmax_idx]
            auc = trapezoid(concentration, time_points)
            
            return {
                'individual_id': individual_params['id'],
                'parameters': individual_params,  # Store parameters for access
                'bcs_classification': bcs_result,
                'adam_absorption': adam_result,
                'concentration': concentration,
                'pk_metrics': {
                    'cmax': cmax,
                    'tmax': tmax,
                    'auc': auc,
                    'bioavailability': adam_result['bioavailability'],
                    'final_absorption': adam_result['fraction_absorbed'][-1]
                },
                'regional_absorption': adam_result['regional_absorbed'],
                'time_points': time_points
            }
            
        except Exception as e:
            print(f"Error processing {individual_params['id']}: {str(e)}")
            return None
    
    def run_population_simulation(self, simulation_time: float = 24.0, time_step: float = 0.1) -> Dict[str, Any]:
        """Run ADAM simulation for entire population"""
        print(f"Running population ADAM simulation for {self.population_size} individuals...")
        
        # Generate population
        population = self.generate_population_parameters()
        
        # Create time points
        time_points = np.arange(0, simulation_time + time_step, time_step)
        
        # Store results
        individual_results = []
        population_concentrations = []
        bcs_classes = []
        
        for i, individual in enumerate(population):
            print(f"Processing individual {i+1}/{self.population_size}", end='\r')
            
            result = self.run_individual_adam_simulation(individual, time_points)
            if result:
                individual_results.append(result)
                population_concentrations.append(result['concentration'])
                bcs_classes.append(result['bcs_classification']['class'])
        
        print(f"\nProcessed {len(individual_results)} individuals successfully")
        
        # Calculate population statistics
        concentrations_array = np.array(population_concentrations)
        
        population_stats = {
            'mean_concentration': np.mean(concentrations_array, axis=0),
            'std_concentration': np.std(concentrations_array, axis=0),
            'median_concentration': np.median(concentrations_array, axis=0),
            'percentile_5': np.percentile(concentrations_array, 5, axis=0),
            'percentile_25': np.percentile(concentrations_array, 25, axis=0),
            'percentile_75': np.percentile(concentrations_array, 75, axis=0),
            'percentile_95': np.percentile(concentrations_array, 95, axis=0),
            'time_points': time_points
        }
        
        # Calculate PK metrics
        cmax_values = [ind['pk_metrics']['cmax'] for ind in individual_results]
        auc_values = [ind['pk_metrics']['auc'] for ind in individual_results]
        tmax_values = [ind['pk_metrics']['tmax'] for ind in individual_results]
        bioavailability_values = [ind['pk_metrics']['bioavailability'] for ind in individual_results]
        
        population_pk = {
            'cmax_stats': {
                'mean': np.mean(cmax_values),
                'std': np.std(cmax_values),
                'median': np.median(cmax_values),
                'cv': np.std(cmax_values) / np.mean(cmax_values) * 100 if np.mean(cmax_values) > 0 else 0,
                'min': np.min(cmax_values),
                'max': np.max(cmax_values)
            },
            'auc_stats': {
                'mean': np.mean(auc_values),
                'std': np.std(auc_values),
                'median': np.median(auc_values),
                'cv': np.std(auc_values) / np.mean(auc_values) * 100 if np.mean(auc_values) > 0 else 0,
                'min': np.min(auc_values),
                'max': np.max(auc_values)
            },
            'tmax_stats': {
                'mean': np.mean(tmax_values),
                'std': np.std(tmax_values),
                'median': np.median(tmax_values),
                'cv': np.std(tmax_values) / np.mean(tmax_values) * 100 if np.mean(tmax_values) > 0 else 0
            },
            'bioavailability_stats': {
                'mean': np.mean(bioavailability_values),
                'std': np.std(bioavailability_values),
                'median': np.median(bioavailability_values),
                'cv': np.std(bioavailability_values) / np.mean(bioavailability_values) * 100 if np.mean(bioavailability_values) > 0 else 0
            }
        }
        
        # BCS classification summary
        bcs_summary = {
            'class_I': bcs_classes.count('I'),
            'class_II': bcs_classes.count('II'),
            'class_III': bcs_classes.count('III'),
            'class_IV': bcs_classes.count('IV')
        }
        
        return {
            'population': population,
            'individual_results': individual_results,
            'population_stats': population_stats,
            'population_pk': population_pk,
            'bcs_summary': bcs_summary,
            'simulation_time': simulation_time,
            'time_step': time_step
        }
    
    def generate_comprehensive_charts(self, results: Dict[str, Any], save_path: str = "../output"):
        """Generate comprehensive ADAM population visualizations"""
        os.makedirs(save_path, exist_ok=True)
        
        time_points = results['population_stats']['time_points']
        mean_conc = results['population_stats']['mean_concentration']
        std_conc = results['population_stats']['std_concentration']
        p5 = results['population_stats']['percentile_5']
        p95 = results['population_stats']['percentile_95']
        p25 = results['population_stats']['percentile_25']
        p75 = results['population_stats']['percentile_75']
        
        individual_results = results['individual_results']
        
        # Create comprehensive visualization
        fig, axes = plt.subplots(4, 3, figsize=(24, 24))
        fig.suptitle('Population ADAM PBPK Modeling Results (100 Individuals)', fontsize=18, fontweight='bold')
        
        # 1. Population concentration-time profile
        ax1 = axes[0, 0]
        ax1.plot(time_points, mean_conc, 'b-', linewidth=3, label='Mean')
        ax1.fill_between(time_points, p5, p95, alpha=0.3, color='blue', label='90% CI')
        ax1.fill_between(time_points, p25, p75, alpha=0.5, color='blue', label='IQR')
        ax1.set_xlabel('Time (hours)')
        ax1.set_ylabel('Plasma Concentration (mg/L)')
        ax1.set_title('Population Concentration-Time Profile')
        ax1.grid(True, alpha=0.3)
        ax1.legend()
        
        # 2. Individual trajectories (sample)
        ax2 = axes[0, 1]
        for i in range(0, min(20, len(individual_results)), 2):  # Show subset
            ax2.plot(time_points, individual_results[i]['concentration'], alpha=0.6, linewidth=1)
        ax2.plot(time_points, mean_conc, 'r-', linewidth=2, label='Mean')
        ax2.set_xlabel('Time (hours)')
        ax2.set_ylabel('Plasma Concentration (mg/L)')
        ax2.set_title('Individual Trajectories (Sample)')
        ax2.grid(True, alpha=0.3)
        ax2.legend()
        
        # 3. Cmax distribution
        ax3 = axes[0, 2]
        cmax_values = [ind['pk_metrics']['cmax'] for ind in individual_results]
        ax3.hist(cmax_values, bins=20, alpha=0.7, color='skyblue', edgecolor='black')
        mean_cmax = np.mean(cmax_values)
        ax3.axvline(mean_cmax, color='red', linestyle='--', linewidth=2, label=f'Mean: {mean_cmax:.1f} mg/L')
        ax3.set_xlabel('Cmax (mg/L)')
        ax3.set_ylabel('Frequency')
        ax3.set_title('Cmax Distribution')
        ax3.legend()
        
        # 4. AUC distribution
        ax4 = axes[1, 0]
        auc_values = [ind['pk_metrics']['auc'] for ind in individual_results]
        ax4.hist(auc_values, bins=20, alpha=0.7, color='lightgreen', edgecolor='black')
        mean_auc = np.mean(auc_values)
        ax4.axvline(mean_auc, color='red', linestyle='--', linewidth=2, label=f'Mean: {mean_auc:.1f} mg·h/L')
        ax4.set_xlabel('AUC (mg·h/L)')
        ax4.set_ylabel('Frequency')
        ax4.set_title('AUC Distribution')
        ax4.legend()
        
        # 5. Bioavailability distribution
        ax5 = axes[1, 1]
        bioavailability_values = [ind['pk_metrics']['bioavailability'] for ind in individual_results]
        ax5.hist(bioavailability_values, bins=20, alpha=0.7, color='lightcoral', edgecolor='black')
        mean_bio = np.mean(bioavailability_values)
        ax5.axvline(mean_bio, color='red', linestyle='--', linewidth=2, label=f'Mean: {mean_bio:.2f}')
        ax5.set_xlabel('Bioavailability')
        ax5.set_ylabel('Frequency')
        ax5.set_title('Bioavailability Distribution')
        ax5.legend()
        
        # 6. Regional absorption analysis
        ax6 = axes[1, 2]
        regional_data = {}
        for region in ['stomach', 'duodenum', 'jejunum', 'ileum', 'colon']:
            regional_values = []
            for ind in individual_results:
                if region in ind['regional_absorption']:
                    total_regional = trapezoid(ind['regional_absorption'][region], time_points)
                    regional_values.append(total_regional)
            if regional_values:
                regional_data[region] = np.mean(regional_values)
        
        if regional_data:
            regions = list(regional_data.keys())
            values = list(regional_data.values())
            colors = plt.cm.Set3(np.linspace(0, 1, len(regions)))
            ax6.bar(regions, values, color=colors)
            ax6.set_xlabel('GI Region')
            ax6.set_ylabel('Mean Absorbed Amount (mg)')
            ax6.set_title('Regional Absorption Distribution')
            ax6.tick_params(axis='x', rotation=45)
        
        # 7. BCS class distribution
        ax7 = axes[2, 0]
        bcs_summary = results['bcs_summary']
        bcs_labels = [f'Class {k.split("_")[-1]}' for k in bcs_summary.keys()]
        bcs_counts = list(bcs_summary.values())
        colors = ['#ff9999', '#66b3ff', '#99ff99', '#ffcc99']
        ax7.pie(bcs_counts, labels=bcs_labels, colors=colors, autopct='%1.1f%%')
        ax7.set_title('BCS Classification Distribution')
        
        # 8. Parameter correlation: Cmax vs Bioavailability
        ax8 = axes[2, 1]
        ax8.scatter(bioavailability_values, cmax_values, alpha=0.6, color='purple')
        ax8.set_xlabel('Bioavailability')
        ax8.set_ylabel('Cmax (mg/L)')
        ax8.set_title('Cmax vs Bioavailability')
        ax8.grid(True, alpha=0.3)
        
        # 9. Age vs PK parameters
        ax9 = axes[2, 2]
        age_values = [ind['parameters']['age'] for ind in individual_results]
        ax9.scatter(age_values, auc_values, alpha=0.6, color='green')
        z = np.polyfit(age_values, auc_values, 1)
        p = np.poly1d(z)
        ax9.plot(age_values, p(age_values), "r--", alpha=0.8)
        ax9.set_xlabel('Age (years)')
        ax9.set_ylabel('AUC (mg·h/L)')
        ax9.set_title('Age vs AUC')
        ax9.grid(True, alpha=0.3)
        
        # 10. Weight vs Cmax
        ax10 = axes[3, 0]
        weight_values = [ind['parameters']['weight'] for ind in individual_results]
        ax10.scatter(weight_values, cmax_values, alpha=0.6, color='orange')
        z = np.polyfit(weight_values, cmax_values, 1)
        p = np.poly1d(z)
        ax10.plot(weight_values, p(weight_values), "r--", alpha=0.8)
        ax10.set_xlabel('Weight (kg)')
        ax10.set_ylabel('Cmax (mg/L)')
        ax10.set_title('Weight vs Cmax')
        ax10.grid(True, alpha=0.3)
        
        # 11. Solubility vs Permeability scatter
        ax11 = axes[3, 1]
        solubility_values = [ind['parameters']['solubility'] for ind in individual_results]
        permeability_values = [ind['parameters']['permeability'] for ind in individual_results]
        ax11.scatter(np.log10(permeability_values), solubility_values, alpha=0.6, color='red')
        ax11.set_xlabel('Log10 Permeability (cm/s)')
        ax11.set_ylabel('Solubility (mg/mL)')
        ax11.set_title('Solubility vs Permeability')
        ax11.grid(True, alpha=0.3)
        
        # 12. Population summary statistics
        ax12 = axes[3, 2]
        ax12.axis('off')
        
        pk_stats = results['population_pk']
        summary_text = f"""Population ADAM PBPK Summary (n=100)

Key PK Parameters:
• Cmax: {pk_stats['cmax_stats']['mean']:.1f} ± {pk_stats['cmax_stats']['std']:.1f} mg/L
  (Range: {pk_stats['cmax_stats']['min']:.1f} - {pk_stats['cmax_stats']['max']:.1f})
  CV: {pk_stats['cmax_stats']['cv']:.1f}%

• AUC: {pk_stats['auc_stats']['mean']:.1f} ± {pk_stats['auc_stats']['std']:.1f} mg·h/L
  (Range: {pk_stats['auc_stats']['min']:.1f} - {pk_stats['auc_stats']['max']:.1f})
  CV: {pk_stats['auc_stats']['cv']:.1f}%

• Bioavailability: {pk_stats['bioavailability_stats']['mean']:.2f} ± {pk_stats['bioavailability_stats']['std']:.2f}
  CV: {pk_stats['bioavailability_stats']['cv']:.1f}%

Population Characteristics:
• Age: {np.mean([ind['parameters']['age'] for ind in individual_results]):.1f} ± {np.std([ind['parameters']['age'] for ind in individual_results]):.1f} years
• Weight: {np.mean([ind['parameters']['weight'] for ind in individual_results]):.1f} ± {np.std([ind['parameters']['weight'] for ind in individual_results]):.1f} kg

BCS Classification:
• Class I: {results['bcs_summary']['class_I']} individuals
• Class II: {results['bcs_summary']['class_II']} individuals
• Class III: {results['bcs_summary']['class_III']} individuals
• Class IV: {results['bcs_summary']['class_IV']} individuals
"""
        
        ax12.text(0.05, 0.95, summary_text, transform=ax12.transAxes, 
                fontsize=10, verticalalignment='top', 
                bbox=dict(boxstyle="round,pad=0.5", facecolor="lightyellow", alpha=0.8))
        
        plt.tight_layout()
        chart_path = os.path.join(save_path, "population_adam_pbpk_comprehensive.png")
        plt.savefig(chart_path, dpi=300, bbox_inches='tight')
        plt.show()
        
        return chart_path
    
    def export_population_data(self, results: Dict[str, Any], save_path: str = "../output"):
        """Export population ADAM data for further analysis"""
        os.makedirs(save_path, exist_ok=True)
        
        individual_results = results['individual_results']
        
        data_list = []
        for ind in individual_results:
            row = {
                'id': ind['individual_id'],
                'age': ind['parameters']['age'],
                'weight': ind['parameters']['weight'],
                'height': ind['parameters']['height'],
                'bmi': ind['parameters']['bmi'],
                'gender': ind['parameters']['gender'],
                'dose': ind['parameters']['dose'],
                'solubility': ind['parameters']['solubility'],
                'permeability': ind['parameters']['permeability'],
                'molecular_weight': ind['parameters']['molecular_weight'],
                'pKa': ind['parameters']['pKa'],
                'logP': ind['parameters']['logP'],
                'particle_size': ind['parameters']['particle_size'],
                'bcs_class': ind['bcs_classification']['class'],
                'cmax': ind['pk_metrics']['cmax'],
                'tmax': ind['pk_metrics']['tmax'],
                'auc': ind['pk_metrics']['auc'],
                'bioavailability': ind['pk_metrics']['bioavailability'],
                'final_absorption': ind['pk_metrics']['final_absorption']
            }
            data_list.append(row)
        
        df = pd.DataFrame(data_list)
        csv_path = os.path.join(save_path, "population_adam_results.csv")
        df.to_csv(csv_path, index=False)
        
        # Export detailed simulation parameters
        params_path = os.path.join(save_path, "population_adam_parameters.json")
        with open(params_path, 'w') as f:
            json.dump({
                'base_drug_params': self.base_drug_params,
                'population_size': self.population_size,
                'simulation_time': results['simulation_time'],
                'time_step': results['time_step'],
                'population_stats': results['population_pk'],
                'bcs_summary': results['bcs_summary']
            }, f, indent=2, default=lambda x: float(x) if isinstance(x, np.ndarray) else x)
        
        print(f"📊 Population ADAM data exported to:")
        print(f"   CSV: {csv_path}")
        print(f"   JSON: {params_path}")
        return csv_path, params_path

def main():
    """Main execution function"""
    print("=" * 80)
    print("POPULATION ADAM PBPK MODELING")
    print("Simulating 100 individuals using ADAM (Advanced Dissolution, Absorption, Metabolism)")
    print("=" * 80)
    
    # Initialize model
    model = PopulationADAMPBPK()
    
    # Run population simulation
    print("\n1. Running population ADAM simulation...")
    results = model.run_population_simulation(simulation_time=24.0, time_step=0.1)
    
    # Display population statistics
    print("\n2. Population ADAM PK Statistics:")
    print("-" * 50)
    pk_stats = results['population_pk']
    
    for param, stats in pk_stats.items():
        print(f"{param.replace('_', ' ').title()}:")
        print(f"  Mean: {stats['mean']:.2f}")
        print(f"  Std: {stats['std']:.2f}")
        if 'cv' in stats:
            print(f"  CV%: {stats['cv']:.1f}%")
        if 'min' in stats:
            print(f"  Range: {stats['min']:.2f} - {stats['max']:.2f}")
        print()
    
    # Display BCS classification summary
    print("3. BCS Classification Summary:")
    print("-" * 30)
    bcs_summary = results['bcs_summary']
    total_individuals = sum(bcs_summary.values())
    for bcs_class, count in bcs_summary.items():
        percentage = (count / total_individuals) * 100
        print(f"Class {bcs_class.split('_')[-1]}: {count} individuals ({percentage:.1f}%)")
    
    # Generate comprehensive visualizations
    print("\n4. Generating comprehensive ADAM population charts...")
    chart_path = model.generate_comprehensive_charts(results)
    print(f"📊 Charts saved to: {chart_path}")
    
    # Export population data
    csv_path, params_path = model.export_population_data(results)
    
    # Display sample individual results
    print("\n5. Sample Individual ADAM Results:")
    print("-" * 50)
    sample_size = min(5, len(results['individual_results']))
    for i in range(sample_size):
        ind = results['individual_results'][i]
        params = next(p for p in results['population'] if p['id'] == ind['individual_id'])
        print(f"Individual {i+1} ({ind['individual_id']}):")
        print(f"  Demographics: {params['age']:.1f} years, {params['weight']:.1f} kg, {params['gender']}")
        print(f"  BCS Class: {ind['bcs_classification']['class']}")
        print(f"  Cmax: {ind['pk_metrics']['cmax']:.2f} mg/L")
        print(f"  AUC: {ind['pk_metrics']['auc']:.1f} mg·h/L")
        print(f"  Bioavailability: {ind['pk_metrics']['bioavailability']:.2f}")
        print()
    
    print("=" * 80)
    print("✅ Population ADAM PBPK Modeling Complete!")
    print("=" * 80)

if __name__ == "__main__":
    main()